from flask import Flask, request, jsonify, send_file
import PyPDF2
import os

app = Flask(__name__)

def encrypt_pdf(input_pdf_path, output_pdf_path, password):
    with open(input_pdf_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        writer = PyPDF2.PdfWriter()

        for page in range(len(reader.pages)):
            writer.add_page(reader.pages[page])

        writer.encrypt(password)

        with open(output_pdf_path, 'wb') as output_file:
            writer.write(output_file)

def decrypt_pdf(input_pdf_path, output_pdf_path, password):
    with open(input_pdf_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)

        if reader.is_encrypted:
            reader.decrypt(password)

        writer = PyPDF2.PdfWriter()

        for page in range(len(reader.pages)):
            writer.add_page(reader.pages[page])

        with open(output_pdf_path, 'wb') as output_file:
            writer.write(output_file)

@app.route('/encrypt', methods=['POST'])
def encrypt():
    file = request.files['file']
    password = request.form['password']
    input_pdf_path = 'temp_input.pdf'
    output_pdf_path = 'temp_encrypted.pdf'

    file.save(input_pdf_path)
    encrypt_pdf(input_pdf_path, output_pdf_path, password)

    return send_file(output_pdf_path, as_attachment=True)

@app.route('/decrypt', methods=['POST'])
def decrypt():
    file = request.files['file']
    password = request.form['password']
    input_pdf_path = 'temp_input.pdf'
    output_pdf_path = 'temp_decrypted.pdf'

    file.save(input_pdf_path)
    decrypt_pdf(input_pdf_path, output_pdf_path, password)

    return send_file(output_pdf_path, as_attachment=True)

if __name__ == '__main__':
    app.run(debug=True) 